"""
统一的任务进度追踪工具
提供详细的进度更新、实时推送和预计时间等功能
"""

import asyncio
import time
from datetime import datetime
from loguru import logger
from typing import Optional, Callable, Any, Dict, List

from app.models.task.task_data import TaskStatus
from app.models.task.progress_data import ProgressData
from app.core.sse_manager import sse_manager


class ProgressTracker:
    """统一的进度追踪器，支持实时推送和详细追踪"""

    def __init__(self, task_id: str, task_name: str = "", storage=None):
        self.task_id = task_id
        self.task_name = task_name
        self.sse_manager = sse_manager
        self.storage = storage

        # 进度数据
        self._progress = 0.0
        self._status = TaskStatus.PENDING
        self._message = ""
        self._details: Dict[str, Any] = {}
        self._start_time = None
        self._end_time = None
        self._steps_completed = 0
        self._total_steps = 0

        # 回调函数列表
        self._callbacks: List[Callable] = []

        # 子任务进度
        self._subtasks: Dict[str, float] = {}

    async def start(self, total_steps: int = 100, message: str = "任务开始..."):
        """开始任务并设置总步数"""
        self._start_time = datetime.now()
        self._total_steps = total_steps
        self._status = TaskStatus.RUNNING
        await self.set_progress(0, message)

    async def set_progress(
        self,
        progress: float,
        message: str = "",
        details: Optional[Dict[str, Any]] = None,
    ):
        """设置进度并推送更新"""
        self._progress = max(0, min(100, progress))  # 限制在0-100之间
        self._message = message or self._message
        if details is not None:
            self._details = details

        # 计算预计剩余时间
        eta = self._calculate_eta()

        progress_data = ProgressData(
            task_id=self.task_id,
            progress=self._progress,
            status=self._status,
            message=self._message,
            details=self._details,
            start_time=self._start_time,
            end_time=self._end_time,
            eta=eta,
            steps_completed=int(self._progress * self._total_steps / 100),
            total_steps=self._total_steps,
            subtasks=self._subtasks.copy(),
            timestamp=datetime.now(),
        )

        # 存储进度
        if self.storage and hasattr(self.storage, "set_progress"):
            try:
                await self.storage.set_progress(
                    self.task_id, progress_data.model_dump()
                )
            except Exception as e:
                logger.error(f"存储进度失败: {e}")

        # SSE实时推送
        if self.sse_manager:
            try:
                await self.sse_manager.broadcast_message(
                    "task_progress", progress_data.model_dump()
                )
                # 为特定任务订阅者推送
                await self.sse_manager.broadcast_message(
                    f"task_{self.task_id}", progress_data.model_dump()
                )
            except Exception as e:
                logger.error(f"SSE推送失败: {e}")

        # 执行回调函数
        for callback in self._callbacks:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(progress_data.model_dump())
                else:
                    callback(progress_data.model_dump())
            except Exception as e:
                logger.error(f"Progress callback error: {e}")

    async def add_step(self, step_name: str, message: str = ""):
        """增加一个完成步骤"""
        self._steps_completed += 1
        progress = (self._steps_completed / self._total_steps) * 100
        step_message = f"步骤 {self._steps_completed}/{self._total_steps}: {step_name}"
        if message:
            step_message += f" - {message}"
        await self.set_progress(progress, step_message)

    async def set_subtask_progress(
        self, subtask_name: str, progress: float, weight: float = 1.0
    ):
        """设置子任务进度（用于复杂任务）"""
        self._subtasks[subtask_name] = progress

        # 计算总体进度（基于子任务权重）
        if self._subtasks:
            total_progress = sum(self._subtasks.values()) / len(self._subtasks)
            await self.set_progress(total_progress, f"子任务进度更新: {subtask_name}")

    async def complete(self, message: str = "任务完成"):
        """标记任务完成"""
        self._end_time = datetime.now()
        self._status = TaskStatus.COMPLETED
        await self.set_progress(100, message)

    async def fail(self, error_message: str, details: Optional[Dict[str, Any]] = None):
        """标记任务失败"""
        self._end_time = datetime.now()
        self._status = TaskStatus.FAILED
        await self.set_progress(self._progress, f"任务失败: {error_message}", details)

    def add_callback(self, callback: Callable):
        """添加进度更新回调"""
        self._callbacks.append(callback)

    def _calculate_eta(self) -> Optional[float]:
        """计算预计剩余时间（秒）"""
        if not self._start_time or self._progress <= 0:
            return None

        elapsed = (datetime.now() - self._start_time).total_seconds()
        if self._progress >= 100:
            return 0

        rate = self._progress / elapsed  # 百分比/秒
        remaining_progress = 100 - self._progress
        return remaining_progress / rate if rate > 0 else None

    @property
    def current_progress(self) -> dict:
        """获取当前进度快照"""
        return {
            "task_id": self.task_id,
            "progress": self._progress,
            "status": self._status.value,
            "message": self._message,
            "details": self._details,
            "start_time": self._start_time.isoformat() if self._start_time else None,
            "end_time": self._end_time.isoformat() if self._end_time else None,
            "eta": self._calculate_eta(),
            "steps_completed": self._steps_completed,
            "total_steps": self._total_steps,
            "subtasks": self._subtasks.copy(),
        }
